//
// Created by hzj on 24-12-24.
//

#include "multi_camera_cooperation/apriltag_location.h"

AprilTagLocationNodeROS::AprilTagLocationNodeROS(ros::NodeHandle &nh){
    printf("[AprilTag Location Node ROS]\n");

    // Load ROS parameters
    nh.param<int>("apriltag_num", apriltag_num, 5);
    nh.param<std::string>("apriltag_location_topic", apriltag_location_topic, "/tag_detections");
    nh.param<std::string>("cam", cam, "CAM_D");
    nh.param<std::string>("config_file_path", config_file_path, "/home/hezijia/catkin_ws/src/multi_camera_cooperation/config/");
    nh.param<int>("minitag_id", minitag_id, 0);

    // ROS publisher and subscriber
    apriltag_detection_sub = nh.subscribe(apriltag_location_topic, 1, &AprilTagLocationNodeROS::apriltags_callback, this);
    // apriltag_location_pub = nh.advertise<geometry_msgs::TransformStamped>("/apriltag_location_estimation", 1);
    apriltag_estimation_pub = nh.advertise<geometry_msgs::TransformStamped>("/apriltag/T_base_to_estimation", 1);

    // Load config
    std::string path = config_file_path;
    LoadConfig(path, cam, apriltag_num);

    correction_matrix(0, 0) = 0.0;
    correction_matrix(0, 1) = 1.0;
    correction_matrix(1, 0) = 1.0;
    correction_matrix(1, 1) = 0.0;
}

void AprilTagLocationNodeROS::apriltags_callback(const apriltag_ros::AprilTagDetectionArray::ConstPtr& msg){
    apriltag_detection_msg = msg;
    
    location_solve();
}

void AprilTagLocationNodeROS::location_solve(){
    apriltag_location_list.clear();
    printf(YELLOW "Detection size: %zu\n" RESET, apriltag_detection_msg->detections.size());
    int detection_size = apriltag_detection_msg->detections.size();

    for (int i = 0; i < detection_size; i++)
    {
        auto detection = apriltag_detection_msg->detections[i];

        double x = detection.pose.pose.pose.position.x;
        double y = detection.pose.pose.pose.position.y;
        double z = detection.pose.pose.pose.position.z;
        Eigen::Vector3d t(x, y, z);
        
        Eigen::Quaterniond q(detection.pose.pose.pose.orientation.w, detection.pose.pose.pose.orientation.x, detection.pose.pose.pose.orientation.y, detection.pose.pose.pose.orientation.z);
        Eigen::Matrix3d R = q.toRotationMatrix();

        Eigen::Matrix4d T = Eigen::Matrix4d::Identity();
        Eigen::Matrix4d T_temp = Eigen::Matrix4d::Identity();
        T.block<3,3>(0,0) = R;
        T.block<3,1>(0,3) = t;

        for(int j = 0; j < apriltag_package_list.size(); j++){
            if(detection.id[0] == apriltag_package_list[j].id){
                T_temp = T;
                T = correction_matrix * ((apriltag_package_list[j].offset * T) * T_base_to_cam.inverse());
                break;
            }
        }

        double weight = 1.0;
        if((detection.id[0] == minitag_id) && (detection_size > 1)){
            weight = 0.0;
        }
        // if(fusion_position != Eigen::Vector3d::Zero()){
        //     Eigen::Vector3d p = T.block<3,1>(0,3);
        //     Eigen::Quaterniond q = Eigen::Quaterniond(T.block<3,3>(0,0));
        //     Eigen::Vector3d p_fusion = fusion_position;
        //     Eigen::Quaterniond q_fusion = fusion_orientation;

        //     double distance = (p - p_fusion).norm(); // meter
        //     double angle = acos(q.dot(q_fusion)) * 180 / M_PI; // degree

        //     weight = max(1.0, 0.1 / distance);
        // }
        // else
        // {
        //     weight = 1.0;
        // }
        
        AprilTagLocation apriltag_location;
        apriltag_location.id = detection.id[0];
        apriltag_location.weight = weight;
        apriltag_location.position = T.block<3,1>(0,3);
        apriltag_location.orientation = Eigen::Quaterniond(T.block<3,3>(0,0));
        apriltag_location.origen_position = T_temp.block<3,1>(0,3);

        apriltag_location_list.emplace_back(apriltag_location);
        
        // // Print the result
        // printf(GREEN "ID: %d\n" RESET, apriltag_location.id);
        // printf(GREEN "Weight: %f\n" RESET, apriltag_location.weight);
        // printf(GREEN "Position: " RESET);
        // for(int j = 0; j < 3; j++){
        //     std::cout << T(j, 3) << " ";
        // }
        // std::cout << std::endl;
        // printf(GREEN "Orientation:\n" RESET);
        // for(int j = 0; j < 3; j++){
        //     for(int k = 0; k < 3; k++){
        //         std::cout << T(j, k) << " ";
        //     }
        //     std::cout << std::endl;
        // }
        
    }

    // result fusion
    // initialize
    fusion_position = Eigen::Vector3d::Zero();
    fusion_orientation = Eigen::Quaterniond::Identity();
    fusion_origin_position = Eigen::Vector3d::Zero();
    double weight_sum = 0;
    for(int i = 0; i < apriltag_location_list.size(); i++){
        fusion_position = fusion_position + apriltag_location_list[i].weight * apriltag_location_list[i].position;
        fusion_orientation = QuaternionsClose_Merge(apriltag_location_list); 
        fusion_origin_position = fusion_origin_position + apriltag_location_list[i].weight * apriltag_location_list[i].origen_position;
        weight_sum += apriltag_location_list[i].weight;   
    }
    fusion_position = fusion_position / weight_sum;
    fusion_orientation.normalize();
    fusion_origin_position = fusion_origin_position / weight_sum;

    printf(GREEN "Fusion Position: " RESET);
    for(int i = 0; i < 3; i++){
        std::cout << fusion_position(i) << " ";
    }
    std::cout << std::endl;
    printf(GREEN "Fusion Orientation:\n" RESET);
    for(int i = 0; i < 3; i++){
        for(int j = 0; j < 3; j++){
            std::cout << fusion_orientation.toRotationMatrix()(i, j) << " ";
        }
        std::cout << std::endl;
    }

    if(fusion_position != Eigen::Vector3d::Zero()){
        // Publish the result
        geometry_msgs::TransformStamped msg;
        msg.header.stamp = ros::Time::now();
        msg.header.frame_id = "base";
        msg.child_frame_id = "estimation";
        msg.transform.translation.x = fusion_position(0);
        msg.transform.translation.y = fusion_position(1);
        msg.transform.translation.z = fusion_origin_position(2);
        msg.transform.rotation.w = fusion_orientation.w();
        msg.transform.rotation.x = fusion_orientation.x();
        msg.transform.rotation.y = fusion_orientation.y();
        msg.transform.rotation.z = fusion_orientation.z();
        apriltag_estimation_pub.publish(msg);
    }
}

void AprilTagLocationNodeROS::LoadConfig(std::string config_file_path, std::string cam_name, int apriltag_num){
    std::string apriltag_config_file_path = config_file_path + "apriltag.yaml";
    std::string cam_config_file_path = config_file_path + cam_name + ".yaml";
    std::ifstream fin(apriltag_config_file_path);
    YAML::Node tag_config = YAML::Load(fin);

    for(int i = 0; i < 4; i++){
        for(int j = 0; j < 4; j++){
            T_base_to_origin(i, j) = tag_config["T_base_origin"][i][j].as<double>();
        }
    }

    for(int i = 0; i < apriltag_num; i++){
        std::string tag_name = "tag_" + std::to_string(i);
        YAML::Node tag = tag_config[tag_name];

        apriltag_package apriltag;
        apriltag.id = tag["id"].as<int>();

        Eigen::Matrix4d offset = Eigen::Matrix4d::Identity();
        for(int j = 0; j < 4; j++){
            for(int k = 0; k < 4; k++){
                offset(j, k) = tag["T_origin_tag"][j][k].as<double>();
            }
        }
        apriltag.offset = T_base_to_origin * offset;

        apriltag_package_list.emplace_back(apriltag);
    }

    std::ifstream fin2(cam_config_file_path);
    YAML::Node cam_config = YAML::Load(fin2);

    for(int i = 0; i < 4; i++){
        for(int j = 0; j < 4; j++){
            T_base_to_cam(i, j) = cam_config["Camera"]["T_base_to_cam"][i][j].as<double>();
        }
    }
}

Eigen::Quaterniond AprilTagLocationNodeROS::QuaternionsClose_Merge(std::vector<AprilTagLocation> &tag){
    Eigen::Quaterniond q;
    double w = 0;
    double x = 0;
    double y = 0;
    double z = 0;
    for(int i = 0; i < tag.size(); i++){
        w += tag[i].orientation.w() * tag[i].weight;
        x += tag[i].orientation.x() * tag[i].weight;
        y += tag[i].orientation.y() * tag[i].weight;
        z += tag[i].orientation.z() * tag[i].weight;
    }
    
    double weight = 0;
    for(int i = 0; i < tag.size(); i++){
        weight += tag[i].weight;
    }
    q.w() = w / weight;
    q.x() = x / weight;
    q.y() = y / weight;
    q.z() = z / weight;

    return q;
}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "apriltag_location_node_ros");
    ros::NodeHandle nh("~");
    ros::Rate rate(30);

    AprilTagLocationNodeROS apriltag_location_node(nh);

    ros::spin();
}